import heapq
import itertools
import os
import datetime
import random
import sys
import time

import cv2
import numpy as np

from libs.action import WebDriverAction
from libs.game import Game
from libs.image import ImageTool
from libs.log import Log
from libs.movement import Movement, KeyCode
from libs.task_helper import TaskHelper
from libs.trade_helper import TradeHelper
from libs.window import AdsWindow

# ===================== 图像处理脚本 =====================
def generate_and_save_cropped_images(driver, save_dir="runtime/puzzle_images", start_x=340, start_y=280, offset=65, matrix_size=4):
    """
    生成固定的坐标矩阵，并根据坐标截取图像，同时将图像存储到内存和本地磁盘。
    :param driver: Selenium 驱动，用于截图
    :param save_dir: 保存图片的本地目录
    :param start_x: 左上角的初始 X 坐标
    :param start_y: 左上角的初始 Y 坐标
    :param offset: 每个格子的偏移量
    :param matrix_size: 矩阵的大小（如 4x4）
    :return: cropped_images（截取的图像列表）和 image_positions（图像编号及坐标）
    """
    cropped_images = []  # 存储图像
    image_positions = []  # 存储图像编号和坐标
    position_matrix = []  # 存储生成的坐标矩阵

    # 确保保存目录存在，如果不存在则创建
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    # 生成矩阵并截取图像
    for i in range(matrix_size):
        row = []
        for j in range(matrix_size):
            # 计算每个格子的坐标
            x = start_x + j * offset
            y = start_y + i * offset
            row.append((x, y))

            # 获取网页截图
            screenshot_data = driver.get_screenshot_as_png()
            screenshot = cv2.imdecode(np.frombuffer(screenshot_data, np.uint8), -1)

            # 截取45x45区域
            cropped_image = screenshot[y:y + 45, x:x + 45]
            image_number = i * matrix_size + j + 1  # 计算图像编号

            # 存储图像到内存
            cropped_images.append(cropped_image)
            image_positions.append(((x, y), image_number))

            # 保存图片到本地磁盘
            image_path = os.path.join(save_dir, f"{image_number}.png")
            cv2.imwrite(image_path, cropped_image)

        position_matrix.append(row)

    # 打印生成的坐标矩阵
    print("\nGenerated Position Matrix:")
    for row in position_matrix:
        print(row)

    return cropped_images, image_positions


def find_target(driver, target_file_name, threshold=0.9):
    # 获取屏幕截图
    screenshot_data = driver.get_screenshot_as_png()
    screenshot = cv2.imdecode(np.frombuffer(screenshot_data, np.uint8), -1)

    # 读取目标图像
    current_dir = os.path.dirname(os.path.abspath(__file__))
    parent_dir = os.path.dirname(current_dir)
    target_path = os.path.join(current_dir, 'runtime/puzzle_images', target_file_name)

    if not os.path.exists(target_path):
        print(f"Error: File {target_file_name} does not exist.")
        return None

    target = cv2.imread(target_path)
    # 模板匹配
    result = cv2.matchTemplate(
        cv2.cvtColor(screenshot, cv2.COLOR_RGB2GRAY),
        cv2.cvtColor(target, cv2.COLOR_RGB2GRAY),
        cv2.TM_CCOEFF_NORMED
    )
    _, max_val, _, max_loc = cv2.minMaxLoc(result)
    if max_val > threshold:
        return max_loc  # 返回小图左上角的位置
    return None

def find_all_images(driver, num_images=16, threshold=0.9):
    found_positions = {}
    for i in range(1, num_images + 1):
        target_file_name = f"{i}.png"
        position = find_target(driver, target_file_name, threshold)
        if position is not None:
            found_positions[i] = position
        else:
            found_positions[i] = None

    # 打印每张图片的 (x, y) 坐标
    # for image_num, position in found_positions.items():
    #     if position is not None:
    #         print(f"Image {image_num}: (x={position[0]}, y={position[1]})")
    #     else:
    #         print(f"Image {image_num}: Not Found")

    # 使用固定的 y 值范围进行分组
    def group_by_fixed_y(positions):
        rows = {1: [], 2: [], 3: [], 4: []}  # 分为四组
        for image_num, position in positions.items():
            if position is None:
                continue
            x, y = position
            if 275 <= y < 340:
                rows[1].append((image_num, position))
            elif 340 <= y < 400:
                rows[2].append((image_num, position))
            elif 400 <= y < 465:
                rows[3].append((image_num, position))
            elif 465 <= y < 530:
                rows[4].append((image_num, position))
        return rows

    # 分组
    grouped_rows = group_by_fixed_y(found_positions)

    # 按每行的 x 坐标排序，并补 0 逻辑
    result_matrix = []
    full_position_matrix = [[None for _ in range(4)] for _ in range(4)]  # 用于记录完整拼图块的位置
    x_ranges = [(330, 395), (395, 460), (460, 525), (525, 590)]  # 每列对应的 x 坐标区间

    for row_num in range(1, 5):
        row = grouped_rows[row_num]
        row.sort(key=lambda item: item[1][0])  # 按 x 坐标排序

        # 根据 x 坐标补 0 的逻辑
        current_row = [0, 0, 0, 0]  # 初始化一行全为 0
        for image_num, position in row:
            x, y = position
            for i, (x_min, x_max) in enumerate(x_ranges):
                if x_min <= x < x_max:
                    current_row[i] = image_num
                    full_position_matrix[row_num - 1][i] = position  # 记录完整拼图块的位置
                    break

        result_matrix.append(current_row)

    return result_matrix, full_position_matrix


# ===================== A* 拼图求解脚本 =====================

# 目标状态
goal_state = [
    [1, 2, 3, 4],
    [5, 6, 7, 8],
    [9, 10, 11, 12],
    [13, 14, 15, 0]
]

# 获取目标状态的位置字典，方便计算曼哈顿距离
goal_positions = {goal_state[i][j]: (i, j) for i in range(4) for j in range(4)}

# 曼哈顿距离的计算函数
def manhattan_distance(state):
    distance = 0
    for i in range(4):
        for j in range(4):
            if state[i][j] != 0:
                x, y = goal_positions[state[i][j]]
                distance += abs(x - i) + abs(y - j)
    return distance

# 获取空白块的位置
def find_blank(state):
    for i in range(4):
        for j in range(4):
            if state[i][j] == 0:
                return i, j

# 生成可能的移动
def generate_moves(state):
    blank_i, blank_j = find_blank(state)
    moves = []
    directions = [(-1, 0), (1, 0), (0, -1), (0, 1)]
    for di, dj in directions:
        new_i, new_j = blank_i + di, blank_j + dj
        if 0 <= new_i < 4 and 0 <= new_j < 4:
            new_state = [row[:] for row in state]
            new_state[blank_i][blank_j], new_state[new_i][new_j] = new_state[new_i][new_j], new_state[blank_i][blank_j]
            moves.append(new_state)
    return moves


# A*算法求解拼图
def solve_puzzle(start_state, time_limit=300):
    # 记录开始时间
    start_time = time.time()

    # 初始化优先队列
    priority_queue = []
    heapq.heappush(priority_queue, (0 + manhattan_distance(start_state), 0, start_state, []))

    # 记录已访问的状态
    visited = set()
    visited.add(tuple(itertools.chain(*start_state)))

    # 初始化步数
    steps = 0

    # 初始化上次打印倒计时的时间
    last_print_time = start_time

    while priority_queue:
        # 检查是否超过时间限制
        current_time = time.time()
        if current_time - start_time > time_limit:
            print(f"Time limit of {time_limit} seconds exceeded. Puzzle not solved.")
            return None

        # 每10秒打印一次进度
        if current_time - last_print_time >= 10:
            elapsed_time = int(current_time - start_time)
            remaining_time = max(0, time_limit - elapsed_time)
            print(f"Elapsed time: {elapsed_time}s, Remaining time: {remaining_time}s, Steps: {steps}")
            last_print_time = current_time

        _, cost, current_state, path = heapq.heappop(priority_queue)

        # 打印当前的状态和路径长度
        # print(f"Step {steps}: Current state: {current_state}, Path length: {len(path)}")

        steps += 1

        # 检查是否到达目标状态
        if current_state == goal_state:
            # 记录结束时间并打印总耗时
            end_time = time.time()
            print(f"Puzzle solved in {steps} steps!")
            print(f"Total time: {end_time - start_time:.4f} seconds")
            return path + [current_state]

        # 生成所有可能的移动
        for move in generate_moves(current_state):
            flattened_move = tuple(itertools.chain(*move))
            if flattened_move not in visited:
                visited.add(flattened_move)
                # 将新的状态推入优先队列
                heapq.heappush(priority_queue,
                               (cost + 1 + manhattan_distance(move), cost + 1, move, path + [current_state]))

    # 如果找不到解决方案
    end_time = time.time()
    print(f"Puzzle not solvable. Total time: {end_time - start_time:.4f} seconds")
    return None

def print_puzzle(state):
    for row in state:
        print(' '.join(f'{num:2}' for num in row))
    print()

# 主逻辑
def process_and_solve(driver):
    # 计算完整拼图状态的坐标矩阵
    print("Step 1: 计算完整拼图状态的坐标矩阵...")
    _, full_position_matrix = find_all_images(driver)  # 获取完整状态的坐标矩阵
    print("\nFull Position Matrix (完整拼图坐标矩阵):")
    for row in full_position_matrix:
        print(row)

    # 插入饼干进入拼图状态
    print("\nStep 2: 插入饼干进入拼图状态...")
    insert_cookie()
    action = WebDriverAction(driver)
    action.escape()

    # 初始化循环变量
    is_solved = False
    max_attempts = 2  # 最大尝试次数
    attempts = 0

    while not is_solved and attempts < max_attempts:
        attempts += 1
        # 获取当前拼图逻辑状态
        print("\nStep 3: 获取拼图开始后的逻辑状态...")
        start_state, _ = find_all_images(driver)  # 只需获取逻辑状态
        print("\nStart State Matrix:")
        for row in start_state:
            print(row)

        origin_state = [
            [1, 2, 3, 4],
            [5, 6, 7, 8],
            [9, 10, 11, 12],
            [13, 14, 15, 16]
        ]
        # 检查是否已达到目标状态
        if start_state == goal_state or start_state == origin_state:
            print("\n拼图已完成，无需进一步操作！")
            is_solved = True
            break

        # 求解拼图逻辑并执行移动
        print("\nStep 4: 求解拼图逻辑并移动图块......(需要等待数分钟）")
        solution_path = solve_puzzle(start_state)

        if solution_path is None:
            print("无法求解拼图，请检查当前状态！")
            break

        # 执行拼图点击
        print("\nStep 5: 执行拼图点击...")
        execute_puzzle_moves(action, solution_path, full_position_matrix)
        print("移动完成，等待操作稳定...")

        # 等待操作完成后再检查状态
        time.sleep(2)

    if not is_solved:
        print("拼图未能完成，达到最大尝试次数或遇到错误！")


# 插入饼干的逻辑
def insert_cookie():
    game.use_backpack_item(cookie_index)
    action.click_by_offset(400, 400)
    action.escape()
    action.escape()
    print("insert_cookie done")

# 点击图块的执行逻辑
def execute_puzzle_moves(action, solution_path, full_position_matrix):
    total_steps = len(solution_path) - 1  # 总步数（去掉初始状态）
    for step_index, step in enumerate(solution_path):
        # 跳过初始状态
        if step_index == 0:
            continue

        # 找到 0 的位置
        zero_position = None
        for i, row in enumerate(step):
            for j, val in enumerate(row):
                if val == 0:
                    zero_position = full_position_matrix[i][j]
                    break
            if zero_position:
                break

        if zero_position:
            # 偏移点击位置 (25, 25)
            offset_position = (zero_position[0] + 25, zero_position[1] + 25)
            print(f"第 {step_index}/{total_steps} 步: 点击位置 {offset_position}")
            action.click_by_offset(offset_position[0], offset_position[1], wait_secons=1)
    time.sleep(3)
    action.handle_dialog()
    action.handle_dialog()
    game.use_backpack_item(cookie_index)
    action.click_by_offset(100, 100)    # 把鼠标移开
    time.sleep(2)



if __name__ == "__main__":
    parser = TaskHelper.default_args_parser('Pixmax Puzzle')
    args = parser.parse_args()
    task_helper = TaskHelper.from_args(args)

    user_id = args.user_id
    log = Log(user_id)

    try:
        driver = AdsWindow(user_id).open(False)
        game = Game(driver, user_id)
        action = WebDriverAction(driver)
        movement = Movement(driver)
        image_tool = ImageTool(driver)
        trade_helper = TradeHelper(driver, user_id)

        game.enter_game()

        cookie_name = 'Cookie Credit'
        cookie_index, cookie_qty = game.get_backpack_index_and_qty(cookie_name)
        log.debug(f'start_cookie_qty={cookie_qty}')
        if cookie_qty < 5:
            task_helper.task_success()
            sys.exit()

        door_loc = None
        for _ in range(3):
            door_loc = image_tool.find_target('carnival/carnival_door.png')
            if door_loc:
                break
            action.goto_spec()
            action.goto_terra()
            movement.keep_move_left(1)
            tent_loc = image_tool.find_target('carnival/carnival_tent.png')
            if tent_loc:
                action.move_to('carnival/carnival_tent.png', max_loop=3)
                action.move_by_offset(200, 200)
                time.sleep(0.5)
                tent_loc = image_tool.find_target('carnival/carnival_tent.png')
                tent_x, tent_y = tent_loc
                action.click_by_offset(tent_x, tent_y)
                action.wait_page_loaded()

        if door_loc is None:
            log.debug('enter carnival failed')
            task_helper.add_task('pixmas_puzzle.py', delay_seconds=0, ignore_running=True, priority=9, retry=True)
            sys.exit()

        tree_loc = None
        for _ in range(3):
            tree_loc = image_tool.find_target('carnival/carnival_tree.png')
            if tree_loc:
                break
            movement.keep_move_up(5)

        movement.keep_move_left(0.5)
        movement.keep_move_up(5)

        movement.keep_move_right(0.5)
        movement.keep_move_up(3)
        movement.keep_move_down(0.3)
        movement.keep_move_right(1.5)
        movement.keep_move_up(3)
        movement.keep_move_down(0.3)
        movement.keep_move_right(2.3)

        movement.send_key_down(KeyCode.UP)
        pl_image = 'carnival/carnival_polar_puzzle.png'
        for _ in range(15):
            ctc_loc = image_tool.find_target(pl_image)
            if ctc_loc is not None:
                break
            time.sleep(0.5)

        movement.send_key_up(KeyCode.UP)
        action.move_to(pl_image, delta_x=0, delta_y=0)
        action.wait_page_loaded()

        movement.keep_move_up(4)
        movement.keep_move_up(4)
        time.sleep(1)
        generate_and_save_cropped_images(driver)

        process_and_solve(driver)

        task_helper.task_success()
        task_helper.add_task('pixmas_puzzle.py', delay_seconds=24 * 60 * 60, ignore_running=True, priority=9)
    except Exception as e:
        task_helper.task_fail()
        task_helper.add_task('pixmas_puzzle.py', delay_seconds=1 * 60, ignore_running=True, priority=9, retry=True)
        raise e



